#include <cuda_runtime.h>
#include <cuda_fp16.h> 
#include "conv2d.h"

extern "C" __global__ void implicit_gemm_v2(mykernelParamType param)
{
    const int ohow = blockIdx.x*16 + threadIdx.x;  // 输出位置合并维度
    const int k = blockIdx.y*16 + threadIdx.y;     // 输出通道维度
    const int n = blockIdx.z;                      // 批次维度
    const int tx = threadIdx.x;
    const int ty = threadIdx.y;
    // 共享内存声明
    __shared__ float sh_input[16][16+1];  
    __shared__ float sh_weight[16][16+1];
    // 分解输出坐标
    const int oh = ohow / param.Ow;
    const int ow = ohow % param.Ow;
    // 输入起始坐标预计算（减少重复计算）
    const int ih_start = oh * param.u - param.p;
    const int iw_start = ow * param.v - param.q;
    float sum = 0.0f;
    // 预计算内存偏移量
    const int input_offset = n * param.c * param.h * param.w;
    const int kernel_offset = k * param.c * param.r * param.s;
    const int crs_total = param.c * param.r * param.s;
    // 分块处理CRS维度（16为分块大小）
    for(int base=0; base<crs_total; base+=16)
    {
        const int crs_idx = base + tx;
        // 权重加载到共享内存
        sh_weight[ty][tx] = (crs_idx < crs_total) ? param.pweight[kernel_offset + crs_idx] : 0.0f;
        // 分解三维索引
        const int c = crs_idx / (param.r * param.s);
        const int residual = crs_idx % (param.r * param.s);
        const int r = residual / param.s;
        const int s = residual % param.s;
        const int ih = ih_start + r;
        const int iw = iw_start + s;
        // 输入加载到共享内存（带边界检查）
        if(ih>=0 && iw>=0 && ih<param.h && iw<param.w)
        {
            const int i_addr = input_offset + c * param.h * param.w + ih * param.w + iw;
            sh_input[ty][tx] = param.pin[i_addr];
        }
        else 
        {
            sh_input[ty][tx] = 0.0f;
        }
        __syncthreads();
        // 矩阵乘累加（展开循环提升ILP）
        #pragma unroll
        for(int i=0; i<16; i++)
        {
            sum += sh_input[ty][i] * sh_weight[i][tx];  // 转置访问模式
        }
        __syncthreads();
    }

    /**** 结果写入 ****/
    if(ohow < param.Oh*param.Ow && k < param.k){
        const int o_addr = n * param.k * param.Oh * param.Ow 
                         + k * param.Oh * param.Ow 
                         + oh * param.Ow 
                         + ow;
        param.pout[o_addr] = sum;
    }
}

void launch_implicit_gemm_v2(unsigned int outh, unsigned int outw, unsigned int k, unsigned int n, mykernelParamType* param) {
    int blockx = (param->Oh * param->Ow + 15) / 16 ;
    int blocky = (param->k + 15) / 16;
    int blockz = param->n;
    int threadx = 16;        
    int thready = 16;
    int threadz = 1;
    dim3 block(threadx, thready, threadz);
    dim3 grid(blockx, blocky, blockz);
    implicit_gemm_v2<<<grid, block>>>(*param);
}